#
import unittest
import torch
from anns.xlstm import MLstmBlock
from anns.xlstm import CausalConv1D
from anns.xlstm import BlockDiagonal

class TXLstm(unittest.TestCase):
    # python -m unittest -v uts.anns.t_xlstm.TXLstm.test_MLstmBlock
    def test_MLstmBlock(self):
        x0 = torch.zeros((1, 5, 64), dtype=torch.float32)
        model = MLstmBlock(x_example=x0, factor=2.0, depth=4, dropout=0.2)
        y = model(x0)
        print(f'MLstmBlock Left: {y.shape}; ')

    # python -m unittest -v uts.anns.t_xlstm.TXLstm.test_CausalConv1D
    def test_CausalConv1D(self):
        x = torch.zeros((1, 5, 16), dtype=torch.float32)
        model = CausalConv1D(in_channels=5, out_channels=6, kernel_size=3)
        y = model(x) # (1, 6, 16)
        print(f'\ny: {y.shape};')

    # python -m unittest -v uts.anns.t_xlstm.TXLstm.test_BlockDiagonal
    def test_BlockDiagonal(self):
        x = torch.zeros((1, 6, 16), dtype=torch.float32)
        model = BlockDiagonal(in_features=16, out_features=8, num_blocks=4)
        y = model(x)
        print(f'BlockDiagonal y: {y.shape};')